{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3YLyiTwPfVCT"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/magenta/ddsp/blob/main/ddsp/colab/demos/timbre_transfer.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n",
        "\n",
        "##### Copyright 2021 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bvp6GWqtfVCW"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 Google LLC. All Rights Reserved.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL"
      },
      "source": [
        "# DDSP Timbre Transfer Demo\n",
        "\n",
        "This notebook is a demo of timbre transfer using DDSP (Differentiable Digital Signal Processing). \n",
        "The model here is trained to generate audio conditioned on a time series of fundamental frequency and loudness. \n",
        "\n",
        "* [DDSP ICLR paper](https://openreview.net/forum?id=B1x1ma4tDr)\n",
        "* [Audio Examples](http://goo.gl/magenta/ddsp-examples) \n",
        "\n",
        "This notebook extracts these features from input audio (either uploaded files, or recorded from the microphone) and resynthesizes with the model. \n",
        "\n",
        "\u003cimg src=\"https://magenta.tensorflow.org/assets/ddsp/ddsp_cat_jamming.png\" alt=\"DDSP Tone Transfer\" width=\"700\"\u003e\n",
        "\n",
        "\n",
        "\n",
        "By default, the notebook will download pre-trained models. You can train a model on your own sounds by using the [Train Autoencoder Colab](https://github.com/magenta/ddsp/blob/main/ddsp/colab/demos/train_autoencoder.ipynb).\n",
        "\n",
        "Have fun! And please feel free to hack this notebook to make your own creative interactions.\n",
        "\n",
        "\n",
        "### Instructions for running:\n",
        "\n",
        "* Make sure to use a GPU runtime, click:  __Runtime \u003e\u003e Change Runtime Type \u003e\u003e GPU__\n",
        "* Press ▶️ on the left of each of the cells\n",
        "* View the code: Double-click any of the cells\n",
        "* Hide the code: Double click the right side of the cell\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6wZde6CBya9k"
      },
      "outputs": [],
      "source": [
        "#@title #Install and Import\n",
        "\n",
        "#@markdown Install ddsp, define some helper functions, and download the model. This transfers a lot of data and _should take a minute or two_.\n",
        "print('Installing from pip package...')\n",
        "!pip install -qU ddsp==1.6.5 \"hmmlearn<=0.2.7\"\n",
        "\n",
        "# Ignore a bunch of deprecation warnings\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "import copy\n",
        "import os\n",
        "import time\n",
        "\n",
        "import crepe\n",
        "import ddsp\n",
        "import ddsp.training\n",
        "from ddsp.colab.colab_utils import (\n",
        "    auto_tune, get_tuning_factor, download, \n",
        "    play, record, specplot, upload, \n",
        "    DEFAULT_SAMPLE_RATE)\n",
        "from ddsp.training.postprocessing import (\n",
        "    detect_notes, fit_quantile_transform\n",
        ")\n",
        "import gin\n",
        "from google.colab import files\n",
        "import librosa\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pickle\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "# Helper Functions\n",
        "sample_rate = DEFAULT_SAMPLE_RATE  # 16000\n",
        "\n",
        "\n",
        "print('Done!')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Go36QW9AS_CD"
      },
      "outputs": [],
      "source": [
        "#@title Record or Upload Audio\n",
        "#@markdown * Either record audio from microphone or upload audio from file (.mp3 or .wav) \n",
        "#@markdown * Audio should be monophonic (single instrument / voice)\n",
        "#@markdown * Extracts fundmanetal frequency (f0) and loudness features. \n",
        "\n",
        "record_or_upload = \"Record\"  #@param [\"Record\", \"Upload (.mp3 or .wav)\"]\n",
        "\n",
        "record_seconds =     5#@param {type:\"number\", min:1, max:10, step:1}\n",
        "\n",
        "if record_or_upload == \"Record\":\n",
        "  audio = record(seconds=record_seconds)\n",
        "else:\n",
        "  # Load audio sample here (.mp3 or .wav3 file)\n",
        "  # Just use the first file.\n",
        "  filenames, audios = upload()\n",
        "  audio = audios[0]\n",
        "if len(audio.shape) == 1:\n",
        "  audio = audio[np.newaxis, :]\n",
        "print('\\nExtracting audio features...')\n",
        "\n",
        "# Plot.\n",
        "specplot(audio)\n",
        "play(audio)\n",
        "\n",
        "# Setup the session.\n",
        "ddsp.spectral_ops.reset_crepe()\n",
        "\n",
        "# Compute features.\n",
        "start_time = time.time()\n",
        "audio_features = ddsp.training.metrics.compute_audio_features(audio)\n",
        "audio_features['loudness_db'] = audio_features['loudness_db'].astype(np.float32)\n",
        "audio_features_mod = None\n",
        "print('Audio features took %.1f seconds' % (time.time() - start_time))\n",
        "\n",
        "\n",
        "TRIM = -15\n",
        "# Plot Features.\n",
        "fig, ax = plt.subplots(nrows=3, \n",
        "                       ncols=1, \n",
        "                       sharex=True,\n",
        "                       figsize=(6, 8))\n",
        "ax[0].plot(audio_features['loudness_db'][:TRIM])\n",
        "ax[0].set_ylabel('loudness_db')\n",
        "\n",
        "ax[1].plot(librosa.hz_to_midi(audio_features['f0_hz'][:TRIM]))\n",
        "ax[1].set_ylabel('f0 [midi]')\n",
        "\n",
        "ax[2].plot(audio_features['f0_confidence'][:TRIM])\n",
        "ax[2].set_ylabel('f0 confidence')\n",
        "_ = ax[2].set_xlabel('Time step [frame]')\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wmSGDWM5yyjm"
      },
      "outputs": [],
      "source": [
        "#@title Load a model\n",
        "#@markdown Run for ever new audio input\n",
        "model = 'Violin' #@param ['Violin', 'Flute', 'Flute2', 'Trumpet', 'Tenor_Saxophone', 'Upload your own (checkpoint folder as .zip)']\n",
        "MODEL = model\n",
        "\n",
        "\n",
        "def find_model_dir(dir_name):\n",
        "  # Iterate through directories until model directory is found\n",
        "  for root, dirs, filenames in os.walk(dir_name):\n",
        "    for filename in filenames:\n",
        "      if filename.endswith(\".gin\") and not filename.startswith(\".\"):\n",
        "        model_dir = root\n",
        "        break\n",
        "  return model_dir \n",
        "\n",
        "if model in ('Violin', 'Flute', 'Flute2', 'Trumpet', 'Tenor_Saxophone'):\n",
        "  # Pretrained models.\n",
        "  PRETRAINED_DIR = '/content/pretrained'\n",
        "  # Copy over from gs:// for faster loading.\n",
        "  !rm -r $PRETRAINED_DIR \u0026\u003e /dev/null\n",
        "  !mkdir $PRETRAINED_DIR \u0026\u003e /dev/null\n",
        "  GCS_CKPT_DIR = 'gs://ddsp/models/timbre_transfer_colab/2021-07-08'\n",
        "  model_dir = os.path.join(GCS_CKPT_DIR, 'solo_%s_ckpt' % model.lower())\n",
        "  \n",
        "  !gsutil cp $model_dir/* $PRETRAINED_DIR \u0026\u003e /dev/null\n",
        "  model_dir = PRETRAINED_DIR\n",
        "  gin_file = os.path.join(model_dir, 'operative_config-0.gin')\n",
        "\n",
        "else:\n",
        "  # User models.\n",
        "  UPLOAD_DIR = '/content/uploaded'\n",
        "  !mkdir $UPLOAD_DIR\n",
        "  uploaded_files = files.upload()\n",
        "\n",
        "  for fnames in uploaded_files.keys():\n",
        "    print(\"Unzipping... {}\".format(fnames))\n",
        "    !unzip -o \"/content/$fnames\" -d $UPLOAD_DIR \u0026\u003e /dev/null\n",
        "  model_dir = find_model_dir(UPLOAD_DIR)\n",
        "  gin_file = os.path.join(model_dir, 'operative_config-0.gin')\n",
        "\n",
        "\n",
        "# Load the dataset statistics.\n",
        "DATASET_STATS = None\n",
        "dataset_stats_file = os.path.join(model_dir, 'dataset_statistics.pkl')\n",
        "print(f'Loading dataset statistics from {dataset_stats_file}')\n",
        "try:\n",
        "  if tf.io.gfile.exists(dataset_stats_file):\n",
        "    with tf.io.gfile.GFile(dataset_stats_file, 'rb') as f:\n",
        "      DATASET_STATS = pickle.load(f)\n",
        "except Exception as err:\n",
        "  print('Loading dataset statistics from pickle failed: {}.'.format(err))\n",
        "\n",
        "\n",
        "# Parse gin config,\n",
        "with gin.unlock_config():\n",
        "  gin.parse_config_file(gin_file, skip_unknown=True)\n",
        "\n",
        "# Assumes only one checkpoint in the folder, 'ckpt-[iter]`.\n",
        "ckpt_files = [f for f in tf.io.gfile.listdir(model_dir) if 'ckpt' in f]\n",
        "ckpt_name = ckpt_files[0].split('.')[0]\n",
        "ckpt = os.path.join(model_dir, ckpt_name)\n",
        "\n",
        "# Ensure dimensions and sampling rates are equal\n",
        "time_steps_train = gin.query_parameter('F0LoudnessPreprocessor.time_steps')\n",
        "n_samples_train = gin.query_parameter('Harmonic.n_samples')\n",
        "hop_size = int(n_samples_train / time_steps_train)\n",
        "\n",
        "time_steps = int(audio.shape[1] / hop_size)\n",
        "n_samples = time_steps * hop_size\n",
        "\n",
        "# print(\"===Trained model===\")\n",
        "# print(\"Time Steps\", time_steps_train)\n",
        "# print(\"Samples\", n_samples_train)\n",
        "# print(\"Hop Size\", hop_size)\n",
        "# print(\"\\n===Resynthesis===\")\n",
        "# print(\"Time Steps\", time_steps)\n",
        "# print(\"Samples\", n_samples)\n",
        "# print('')\n",
        "\n",
        "gin_params = [\n",
        "    'Harmonic.n_samples = {}'.format(n_samples),\n",
        "    'FilteredNoise.n_samples = {}'.format(n_samples),\n",
        "    'F0LoudnessPreprocessor.time_steps = {}'.format(time_steps),\n",
        "    'oscillator_bank.use_angular_cumsum = True',  # Avoids cumsum accumulation errors.\n",
        "]\n",
        "\n",
        "with gin.unlock_config():\n",
        "  gin.parse_config(gin_params)\n",
        "\n",
        "\n",
        "# Trim all input vectors to correct lengths \n",
        "for key in ['f0_hz', 'f0_confidence', 'loudness_db']:\n",
        "  audio_features[key] = audio_features[key][:time_steps]\n",
        "audio_features['audio'] = audio_features['audio'][:, :n_samples]\n",
        "\n",
        "\n",
        "# Set up the model just to predict audio given new conditioning\n",
        "model = ddsp.training.models.Autoencoder()\n",
        "model.restore(ckpt)\n",
        "\n",
        "# Build model by running a batch through it.\n",
        "start_time = time.time()\n",
        "_ = model(audio_features, training=False)\n",
        "print('Restoring model took %.1f seconds' % (time.time() - start_time))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "uQFUlIJ_5r36"
      },
      "outputs": [],
      "source": [
        "#@title Modify conditioning\n",
        "\n",
        "#@markdown These models were not explicitly trained to perform timbre transfer, so they may sound unnatural if the incoming loudness and frequencies are very different then the training data (which will always be somewhat true). \n",
        "\n",
        "\n",
        "#@markdown ## Note Detection\n",
        "\n",
        "#@markdown You can leave this at 1.0 for most cases\n",
        "threshold = 1 #@param {type:\"slider\", min: 0.0, max:2.0, step:0.01}\n",
        "\n",
        "\n",
        "#@markdown ## Automatic\n",
        "\n",
        "ADJUST = True #@param{type:\"boolean\"}\n",
        "\n",
        "#@markdown Quiet parts without notes detected (dB)\n",
        "quiet = 20 #@param {type:\"slider\", min: 0, max:60, step:1}\n",
        "\n",
        "#@markdown Force pitch to nearest note (amount)\n",
        "autotune = 0 #@param {type:\"slider\", min: 0.0, max:1.0, step:0.1}\n",
        "\n",
        "#@markdown ## Manual\n",
        "\n",
        "\n",
        "#@markdown Shift the pitch (octaves)\n",
        "pitch_shift =  0 #@param {type:\"slider\", min:-2, max:2, step:1}\n",
        "\n",
        "#@markdown Adjust the overall loudness (dB)\n",
        "loudness_shift = 0 #@param {type:\"slider\", min:-20, max:20, step:1}\n",
        "\n",
        "\n",
        "audio_features_mod = {k: v.copy() for k, v in audio_features.items()}\n",
        "\n",
        "\n",
        "## Helper functions.\n",
        "def shift_ld(audio_features, ld_shift=0.0):\n",
        "  \"\"\"Shift loudness by a number of ocatves.\"\"\"\n",
        "  audio_features['loudness_db'] += ld_shift\n",
        "  return audio_features\n",
        "\n",
        "\n",
        "def shift_f0(audio_features, pitch_shift=0.0):\n",
        "  \"\"\"Shift f0 by a number of ocatves.\"\"\"\n",
        "  audio_features['f0_hz'] *= 2.0 ** (pitch_shift)\n",
        "  audio_features['f0_hz'] = np.clip(audio_features['f0_hz'], \n",
        "                                    0.0, \n",
        "                                    librosa.midi_to_hz(110.0))\n",
        "  return audio_features\n",
        "\n",
        "\n",
        "mask_on = None\n",
        "\n",
        "if ADJUST and DATASET_STATS is not None:\n",
        "  # Detect sections that are \"on\".\n",
        "  mask_on, note_on_value = detect_notes(audio_features['loudness_db'],\n",
        "                                        audio_features['f0_confidence'],\n",
        "                                        threshold)\n",
        "\n",
        "  if np.any(mask_on):\n",
        "    # Shift the pitch register.\n",
        "    target_mean_pitch = DATASET_STATS['mean_pitch']\n",
        "    pitch = ddsp.core.hz_to_midi(audio_features['f0_hz'])\n",
        "    mean_pitch = np.mean(pitch[mask_on])\n",
        "    p_diff = target_mean_pitch - mean_pitch\n",
        "    p_diff_octave = p_diff / 12.0\n",
        "    round_fn = np.floor if p_diff_octave \u003e 1.5 else np.ceil\n",
        "    p_diff_octave = round_fn(p_diff_octave)\n",
        "    audio_features_mod = shift_f0(audio_features_mod, p_diff_octave)\n",
        "\n",
        "\n",
        "    # Quantile shift the note_on parts.\n",
        "    _, loudness_norm = fit_quantile_transform(\n",
        "        audio_features['loudness_db'],\n",
        "        mask_on,\n",
        "        inv_quantile=DATASET_STATS['quantile_transform'])\n",
        "\n",
        "    # Turn down the note_off parts.\n",
        "    mask_off = np.logical_not(mask_on)\n",
        "    loudness_norm[mask_off] -=  quiet * (1.0 - note_on_value[mask_off][:, np.newaxis])\n",
        "    loudness_norm = np.reshape(loudness_norm, audio_features['loudness_db'].shape)\n",
        "    \n",
        "    audio_features_mod['loudness_db'] = loudness_norm \n",
        "\n",
        "    # Auto-tune.\n",
        "    if autotune:\n",
        "      f0_midi = np.array(ddsp.core.hz_to_midi(audio_features_mod['f0_hz']))\n",
        "      tuning_factor = get_tuning_factor(f0_midi, audio_features_mod['f0_confidence'], mask_on)\n",
        "      f0_midi_at = auto_tune(f0_midi, tuning_factor, mask_on, amount=autotune)\n",
        "      audio_features_mod['f0_hz'] = ddsp.core.midi_to_hz(f0_midi_at)\n",
        "\n",
        "  else:\n",
        "    print('\\nSkipping auto-adjust (no notes detected or ADJUST box empty).')\n",
        "\n",
        "else:\n",
        "  print('\\nSkipping auto-adujst (box not checked or no dataset statistics found).')\n",
        "\n",
        "# Manual Shifts.\n",
        "audio_features_mod = shift_ld(audio_features_mod, loudness_shift)\n",
        "audio_features_mod = shift_f0(audio_features_mod, pitch_shift)\n",
        "\n",
        "\n",
        "\n",
        "# Plot Features.\n",
        "has_mask = int(mask_on is not None)\n",
        "n_plots = 3 if has_mask else 2 \n",
        "fig, axes = plt.subplots(nrows=n_plots, \n",
        "                      ncols=1, \n",
        "                      sharex=True,\n",
        "                      figsize=(2*n_plots, 8))\n",
        "\n",
        "if has_mask:\n",
        "  ax = axes[0]\n",
        "  ax.plot(np.ones_like(mask_on[:TRIM]) * threshold, 'k:')\n",
        "  ax.plot(note_on_value[:TRIM])\n",
        "  ax.plot(mask_on[:TRIM])\n",
        "  ax.set_ylabel('Note-on Mask')\n",
        "  ax.set_xlabel('Time step [frame]')\n",
        "  ax.legend(['Threshold', 'Likelihood','Mask'])\n",
        "\n",
        "ax = axes[0 + has_mask]\n",
        "ax.plot(audio_features['loudness_db'][:TRIM])\n",
        "ax.plot(audio_features_mod['loudness_db'][:TRIM])\n",
        "ax.set_ylabel('loudness_db')\n",
        "ax.legend(['Original','Adjusted'])\n",
        "\n",
        "ax = axes[1 + has_mask]\n",
        "ax.plot(librosa.hz_to_midi(audio_features['f0_hz'][:TRIM]))\n",
        "ax.plot(librosa.hz_to_midi(audio_features_mod['f0_hz'][:TRIM]))\n",
        "ax.set_ylabel('f0 [midi]')\n",
        "_ = ax.legend(['Original','Adjusted'])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SLwg1WkHCXQO"
      },
      "outputs": [],
      "source": [
        "#@title #Resynthesize Audio\n",
        "\n",
        "af = audio_features if audio_features_mod is None else audio_features_mod\n",
        "\n",
        "# Run a batch of predictions.\n",
        "start_time = time.time()\n",
        "outputs = model(af, training=False)\n",
        "audio_gen = model.get_audio_from_outputs(outputs)\n",
        "print('Prediction took %.1f seconds' % (time.time() - start_time))\n",
        "\n",
        "# Plot\n",
        "print('Original')\n",
        "play(audio)\n",
        "\n",
        "print('Resynthesis')\n",
        "play(audio_gen)\n",
        "\n",
        "specplot(audio)\n",
        "plt.title(\"Original\")\n",
        "\n",
        "specplot(audio_gen)\n",
        "_ = plt.title(\"Resynthesis\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "3YLyiTwPfVCT"
      ],
      "last_runtime": {},
      "name": "timbre_transfer.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
